#define _GNU_SOURCE 1
#include <arpa/inet.h>
#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <poll.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

#define USER_LIMIT 5
#define BUFFER_SIZE 64
#define FD_LIMIT 65535

/* 客户数据：客户端 socket 地址、待写到客户端的数据的位置、从客户端读入的数据 */
struct client_data {
  struct sockaddr_in address;
  char *write_buf;
  char buf[BUFFER_SIZE];
};

/* 将文件描述符设置成非阻塞 */
int setnonblocking(int fd) {
  int old_option = fcntl(fd, F_GETFL);
  int new_option = old_option | O_NONBLOCK;
  fcntl(fd, F_SETFL, new_option);
  return old_option;
}

int main(int argc, char *argv[]) {
  const char *ip = argv[1];
  int port = atoi(argv[2]);

  int ret = 0;
  struct sockaddr_in address;
  bzero(&address, sizeof(address));
  address.sin_family = AF_INET;
  inet_pton(AF_INET, ip, &address.sin_addr);
  address.sin_port = htons(port);

  int listenfd = socket(PF_INET, SOCK_STREAM, 0);
  assert(listenfd >= 0);

  ret = bind(listenfd, (struct sockaddr *)&address, sizeof(address));
  assert(ret != -1);

  ret = listen(listenfd, 5);
  assert(ret != -1);

  /* 创建 user 数组，分配 FD_LIMIT 个 client_data 对象。可以预期：每个可能的
   * socket 连接都可以获得一个这样的对象 并且 socket
   * 的值可以直接用来索引（作为数组的下标）socket 连接对应的 client_data
   * 对象，这是将 socket 和客户数据关联的简单而高效的数据 */
  struct client_data *users =
      (struct client_data *)malloc(sizeof(struct client_data) * FD_LIMIT);
  /* 尽管我们分配了足够多的 client_data 对象，但为了提高 poll
   * 的性能，仍然有必要限制用户的数量 */
  struct pollfd fds[USER_LIMIT + 1];
  int user_count = 0;
  for (int i = 1; i <= user_count; ++i) {
    fds[i].fd = -1;
    fds[i].events = 0;
  }
  fds[0].fd = listenfd;
  fds[0].events = POLLIN | POLLERR;
  fds[0].revents = 0;

  while (1) {
    ret = poll(fds, user_count + 1, -1);
    if (ret < 0) {
      printf("poll is failed.\n");
      break;
    }
    for (int i = 0; i < user_count + 1; ++i) {
      if ((fds[i].fd == listenfd) && (fds[i].revents & POLLIN)) {
        struct sockaddr_in client_address;
        socklen_t client_addrlength = sizeof(client_address);
        int connfd = accept(listenfd, (struct sockaddr *)&client_address,
                            &client_addrlength);
        if (connfd < 0) {
          printf("errno is: %d\n", errno);
          continue;
        }
        /* 如果请求太多，则关闭新到的链接 */
        if (user_count >= USER_LIMIT) {
          const char *info = "too many users\n";
          printf("%s\n", info);
          send(connfd, info, strlen(info), 0);
          close(connfd);
          continue;
        }
        /* 对于新的链接，同时修改 fds 和
         users 数组。前文已经提到，users[connfd] 对应于新连接文件描述符 connfd
         的客户端数据 */
        user_count++;
        users[connfd].address = client_address;
        setnonblocking(connfd);
        fds[user_count].fd = connfd;
        fds[user_count].events = POLLIN | POLLRDHUP | POLLERR;
        fds[user_count].revents = 0;
        printf("comes a new user, now have %d users\n", user_count);
      } else if (fds[i].revents & POLLERR) {
        printf("get an error from %d\n", fds[i].fd);
        char errors[100];
        memset(errors, 0x0, 100);
        socklen_t length = sizeof(errors);
        if (getsockopt(fds[i].fd, SOL_SOCKET, SO_ERROR, &errors, &length) < 0) {
          printf("get socket option failed\n");
        }
        continue;
      } else if (fds[i].revents & POLLRDHUP) {
        /* 如果客户端关闭链接，则服务器也关闭对应的链接，并将用户总数 - 1 */
        users[fds[i].fd] = users[fds[user_count].fd];
        close(fds[i].fd);
        fds[i] = fds[user_count];
        i--;
        user_count--;
        printf("a client left\n");
      } else if (fds[i].revents & POLLIN) {
        int connfd = fds[i].fd;
        memset(users[connfd].buf, 0x0, BUFFER_SIZE);
        ret = recv(connfd, users[connfd].buf, BUFFER_SIZE - 1, 0);
        printf("get %d bytes of client data %s from %d\n", ret,
               users[connfd].buf, connfd);
        if (ret < 0) {
          /* 如果读操作出错，则关闭链接 */
          if (errno != EAGAIN) {
            close(connfd);
            users[fds[i].fd] = users[fds[user_count].fd];
            fds[i] = fds[user_count];
            i--;
            user_count--;
          }
        } else if (ret == 0) {
        } else {
          /* 如果接收到客户端数据，则通知其他 socket 链接准备写数据 */
          for (int j = 1; j <= user_count; ++j) {
            if (fds[j].fd == connfd) {
              continue;
            }
            fds[j].events |= ~POLLIN;
            fds[j].events |= POLLOUT;
            users[fds[j].fd].write_buf = users[connfd].buf;
          }
        }
      } else if (fds[i].revents & POLLOUT) {
        int connfd = fds[i].fd;
        if (!users[connfd].write_buf) {
          continue;
        }
        ret = send(connfd, users[connfd].write_buf,
                   strlen(users[connfd].write_buf), 0);
        users[connfd].write_buf = NULL;
        /* 写完数据后需要重新注册 fds[i] 上的可读事件 */
        fds[i].events |= ~POLLOUT;
        fds[i].events |= POLLIN;
      }
    }
  }

  free(users);
  close(listenfd);
  return 0;
}